import re
from docx import Document
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.enum.table import WD_CELL_VERTICAL_ALIGNMENT
from docx.shared import Pt


def set_cell_margins(cell, **kwargs):
    """
    设置单元格边距，kwargs 可包含 top, start, bottom, end，单位为 dxa（1 dxa = 1/20 点）
    """
    tc = cell._tc
    tcPr = tc.get_or_add_tcPr()
    tcMar = tcPr.find(qn('w:tcMar'))
    if tcMar is None:
        tcMar = OxmlElement('w:tcMar')
        tcPr.append(tcMar)
    for m in ['top', 'start', 'bottom', 'end']:
        if m in kwargs:
            node = tcMar.find(qn('w:' + m))
            if node is None:
                node = OxmlElement('w:' + m)
                tcMar.append(node)
            node.set(qn('w:w'), str(kwargs[m]))
            node.set(qn('w:type'), 'dxa')


def get_or_add_tblPr(tbl):
    tblPr = tbl.find(qn('w:tblPr'))
    if tblPr is None:
        tblPr = OxmlElement('w:tblPr')
        tbl.insert(0, tblPr)
    return tblPr


def parse_sql(sql):
    """
    解析 SQL 建表语句，提取表名和字段信息。
    每个字段返回：字段名、数据类型、是否为空、索引类型、字段注释内容
    其中：
      - 字段名使用反引号包裹
      - “是否为空”使用 YES 或 NO 表示（包含 NOT NULL 则为 NO，否则 YES）
      - 支持解析表级约束中的 PRIMARY KEY、UNIQUE KEY 和 KEY
    """
    sql = " ".join(sql.split())
    m = re.search(r'CREATE TABLE\s+`?(\w+)`?\s*\((.*)\)\s*', sql, re.IGNORECASE | re.DOTALL)
    if not m:
        raise ValueError("SQL 格式不正确")
    table_name = m.group(1)
    inner = m.group(2).strip()

    # 按逗号分隔，避免 DECIMAL(10,2) 中的逗号干扰
    parts = re.split(r',(?![^(]*\))', inner)
    col_defs = []
    constraints = []
    for part in parts:
        part = part.strip()
        # 如果以反引号开头认为是字段定义，否则为表约束
        if part.startswith("`"):
            col_defs.append(part)
        else:
            constraints.append(part)

    # 解析表约束，建立字段与索引类型的映射
    index_map = {}
    for cons in constraints:
        cons_up = cons.upper()
        if cons_up.startswith("PRIMARY KEY"):
            mcols = re.search(r'\((.*?)\)', cons)
            if mcols:
                cols_str = mcols.group(1)
                cols = [c.strip().strip('`') for c in cols_str.split(',')]
                for col in cols:
                    index_map[col] = "PRIMARY KEY"
        elif cons_up.startswith("UNIQUE KEY"):
            mcols = re.search(r'\((.*?)\)', cons)
            if mcols:
                cols_str = mcols.group(1)
                cols = [c.strip().strip('`') for c in cols_str.split(',')]
                for col in cols:
                    index_map[col] = "UNIQUE"
        elif cons_up.startswith("KEY"):
            mcols = re.search(r'\((.*?)\)', cons)
            if mcols:
                cols_str = mcols.group(1)
                cols = [c.strip().strip('`') for c in cols_str.split(',')]
                for col in cols:
                    # 如果未标记为 PRIMARY KEY 或 UNIQUE，则标记为 KEY
                    if col not in index_map:
                        index_map[col] = "KEY"

    parsed_columns = []
    for col in col_defs:
        # 提取字段名（在第一对反引号内）
        mcol = re.search(r'`([^`]+)`', col)
        if not mcol:
            continue
        field_name = mcol.group(1)
        # 提取数据类型：紧跟字段名之后的部分
        mtype = re.search(r'`[^`]+`\s+([^\s]+)', col)
        data_type = mtype.group(1) if mtype else ""
        # 是否为空：包含 NOT NULL 则为 NO，否则 YES
        nullable = "NO" if "NOT NULL" in col.upper() else "YES"
        # 提取 COMMENT 部分（假设用单引号包裹）
        mcomment = re.search(r"COMMENT\s+'(.*?)'", col, re.IGNORECASE)
        comment = mcomment.group(1) if mcomment else ""
        # 根据索引映射判断索引类型
        index_type = index_map.get(field_name, "")
        parsed_columns.append((field_name, data_type, nullable, index_type, comment))

    return table_name, parsed_columns


def add_three_line_table_border(table):
    """
    为整个表格设置上边框和下边框，形成三线表样式（表头下方单独添加分隔线）
    """
    tbl = table._element
    tblPr = get_or_add_tblPr(tbl)
    tblBorders = OxmlElement('w:tblBorders')

    # 顶边
    top = OxmlElement('w:top')
    top.set(qn('w:val'), 'single')
    top.set(qn('w:sz'), '4')
    top.set(qn('w:space'), '0')
    top.set(qn('w:color'), 'auto')
    tblBorders.append(top)

    # 底边
    bottom = OxmlElement('w:bottom')
    bottom.set(qn('w:val'), 'single')
    bottom.set(qn('w:sz'), '4')
    bottom.set(qn('w:space'), '0')
    bottom.set(qn('w:color'), 'auto')
    tblBorders.append(bottom)

    tblPr.append(tblBorders)


def set_runs_font(paragraph, font_name="Times New Roman", font_size=Pt(11)):
    """
    设置段落中所有 run 的字体为指定字体及字号
    """
    for run in paragraph.runs:
        run.font.name = font_name
        run.font.size = font_size


def create_word_table_from_sql(sql):
    """
    根据 SQL 建表语句生成包含字段名、数据类型、是否为空、索引类型、字段注释内容的三线表，
    表头采用宋体（斜体），数据部分采用 Times New Roman，
    所有内容水平和垂直居中且无额外边距
    """
    table_name, columns = parse_sql(sql)
    document = Document()
    # 设置文档默认字体为 Times New Roman（数据部分）
    style = document.styles['Normal']
    style.font.name = "Times New Roman"

    document.add_heading(f"表: {table_name}", level=1)

    headers = ["字段名", "数据类型", "是否为空", "索引类型", "字段注释内容"]
    table = document.add_table(rows=1, cols=len(headers))
    hdr_cells = table.rows[0].cells
    for idx, header in enumerate(headers):
        hdr_cells[idx].text = header
        for para in hdr_cells[idx].paragraphs:
            para.alignment = WD_ALIGN_PARAGRAPH.CENTER
            para.paragraph_format.space_before = 0
            para.paragraph_format.space_after = 0
            # 表头使用宋体，设置斜体
            set_runs_font(para, font_name="宋体")
            for run in para.runs:
                run.italic = True
        hdr_cells[idx].vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER
        set_cell_margins(hdr_cells[idx], top=0, bottom=0, start=0, end=0)

    for field_name, data_type, nullable, index_type, comment in columns:
        row_cells = table.add_row().cells
        row_cells[0].text = field_name
        row_cells[1].text = data_type
        row_cells[2].text = nullable
        row_cells[3].text = index_type
        row_cells[4].text = comment
        for cell in row_cells:
            for para in cell.paragraphs:
                para.alignment = WD_ALIGN_PARAGRAPH.CENTER
                para.paragraph_format.space_before = 0
                para.paragraph_format.space_after = 0
                # 数据部分使用 Times New Roman
                set_runs_font(para, font_name="Times New Roman")
            cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER
            set_cell_margins(cell, top=0, bottom=0, start=0, end=0)

    add_three_line_table_border(table)

    # 为表头单元格添加下边框，形成表头下方分隔线
    for cell in table.rows[0].cells:
        tc = cell._tc
        tcPr = tc.get_or_add_tcPr()
        cell_borders = OxmlElement('w:tcBorders')
        bottom = OxmlElement('w:bottom')
        bottom.set(qn('w:val'), 'single')
        bottom.set(qn('w:sz'), '4')
        bottom.set(qn('w:space'), '0')
        bottom.set(qn('w:color'), 'auto')
        cell_borders.append(bottom)
        tcPr.append(cell_borders)

    document.save("三线表.docx")
    print("Word 文档已生成：三线表.docx")


if __name__ == "__main__":
    sql_statement = """
CREATE TABLE `sys_user` (
  `user_id` bigint NOT NULL AUTO_INCREMENT COMMENT '用户ID',
  `dept_id` bigint DEFAULT NULL COMMENT '部门ID',
  `user_name` varchar(30) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL COMMENT '用户账号',
  `nick_name` varchar(30) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL COMMENT '用户昵称',
  `user_type` varchar(2) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '00' COMMENT '用户类型（00系统用户）',
  `email` varchar(50) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '' COMMENT '用户邮箱',
  `phonenumber` varchar(11) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '' COMMENT '手机号码',
  `sex` char(1) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '0' COMMENT '用户性别（0男 1女 2未知）',
  `avatar` varchar(100) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '' COMMENT '头像地址',
  `password` varchar(100) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '' COMMENT '密码',
  `status` char(1) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '0' COMMENT '帐号状态（0正常 1停用）',
  `del_flag` char(1) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '0' COMMENT '删除标志（0代表存在 2代表删除）',
  `login_ip` varchar(128) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '' COMMENT '最后登录IP',
  `login_date` datetime DEFAULT NULL COMMENT '最后登录时间',
  `create_by` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '' COMMENT '创建者',
  `create_time` datetime DEFAULT NULL COMMENT '创建时间',
  `update_by` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT '' COMMENT '更新者',
  `update_time` datetime DEFAULT NULL COMMENT '更新时间',
  `remark` varchar(500) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '备注',
  PRIMARY KEY (`user_id`) USING BTREE
) ENGINE=InnoDB AUTO_INCREMENT=101 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci ROW_FORMAT=DYNAMIC COMMENT='后台用户信息表'
    """
    create_word_table_from_sql(sql_statement)
